#include "device.h"

#include "log.h"
#include "string.h"
#include "time.h"

using namespace windows::kernal;

VOID Device::CancelRoutine(_DEVICE_OBJECT *device_object, _IRP *irp) {
  ENTER();
  INFO("Device::CancelRoutine device_object=%p, irp=%p\n", device_object, irp);
  irp->IoStatus.Status = STATUS_CANCELLED;
  irp->IoStatus.Information = 0;
  IoCompleteRequest(irp, IO_NO_INCREMENT);
  IoReleaseCancelSpinLock(irp->CancelIrql);
  LEAVE();
}

Device::Device(PDRIVER_OBJECT driver_object)
  : driver_object_(driver_object), device_object_(NULL) {
  ENTER();
  RtlZeroMemory(device_name_, sizeof(device_name_));
  RtlZeroMemory(symbolic_link_name_, sizeof(symbolic_link_name_));
  InitializeListHead(&pending_irp_list_.node);
  pending_irp_list_.irp = NULL;
  LEAVE();
}

Device::~Device() {
  ENTER();
  Destroy();
  LEAVE();
}

NTSTATUS Device::Create(LPCWSTR device_name, LPCWSTR symbolic_link_name) {
  PAGED_CODE();
  ENTER();

  NTSTATUS status;
  if (!device_name || !wcslen(device_name)) 
    return STATUS_INVALID_PARAMETER_1;
  if (!symbolic_link_name || !wcslen(symbolic_link_name))
    return STATUS_INVALID_PARAMETER_2;

  Destroy();
  status = CreateDevice(device_name);
  if (NT_SUCCESS(status)) {
    status = CreateSymbolicLink(symbolic_link_name);
    if (NT_SUCCESS(status)) {
      LEAVE();
      return STATUS_SUCCESS;
    } else {
      ERR("CreateSymbolicLink(%ws) status=%#X\n", symbolic_link_name, status);
    }
  } else {
    ERR("CreateDevice(%ws) status=%#X\n", device_name, status);
  }
  Destroy();

  LEAVE();
  return STATUS_UNSUCCESSFUL;
}

VOID Device::Destroy() { 
  ENTER();
  PAGED_CODE();

  if (device_object_) {
    if (UNICODE_NULL != symbolic_link_name_[0]) {
      UNICODE_STRING symbolic_link_name;
      RtlInitUnicodeString(&symbolic_link_name, symbolic_link_name_);
      NTSTATUS status = IoDeleteSymbolicLink(&symbolic_link_name);
      if (!NT_SUCCESS(status)) {
        ERR("IoDeleteSymbolicLink fail, symbolic_link_name_=%ws, status=%#X\n", symbolic_link_name_, status);
      } else {
        INFO("IoDeleteSymbolicLink success, symbolic_link_name_=%ws)\n", symbolic_link_name_);
      }
      symbolic_link_name_[0] = UNICODE_NULL;
    }

    INFO("IoDeleteDevice device_object_=%p)\n", device_object_);
    IoDeleteDevice(device_object_);
    device_object_ = NULL;
    device_name_[0] = UNICODE_NULL;
    driver_object_ = NULL;
  }

  LEAVE();
}

NTSTATUS Device::CreateDevice(LPCWSTR device_name) {
  PAGED_CODE();
  ENTER();

  NTSTATUS status;
  UNICODE_STRING name;
  SIZE_T size;

  size = sizeof(device_name_);
  status = FormatDeviceName(device_name, device_name_, size);
  if (!NT_SUCCESS(status)) {
    ERR("FormatDeviceName(%ws) fail, status=%#X\n", device_name, status);
    return status;
  }
  RtlInitUnicodeString(&name, device_name_);
  status = IoCreateDevice(driver_object_, 0, &name, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, TRUE, &device_object_);
  if (!NT_SUCCESS(status)) {
    ERR("IoCreateDevice(%wZ) status=%#X\n", &name, status);
    LEAVE();
    return status;
  }
  device_object_->Flags |= DO_BUFFERED_IO;
  device_object_->Flags &= ~DO_DEVICE_INITIALIZING;
  INFO("IoCreateDevice Success, device_name=%ws, device_object=%p\n", device_name_, device_object_);

  LEAVE();
  return status;
}

NTSTATUS Device::CreateSymbolicLink(LPCWSTR symbolic_link_name) {
  ENTER();
  PAGED_CODE();

  NTSTATUS status;
  UNICODE_STRING device_name, name;
  SIZE_T size;

  size = sizeof(symbolic_link_name_);
  status = FormatSymbolicLinkName(symbolic_link_name, symbolic_link_name_, size);
  if (!NT_SUCCESS(status)) {
    ERR("Device::CreateSymbolicLink FormatSymbolicLinkName status=%#X\n", status);
    return status;
  }
  RtlInitUnicodeString(&device_name, device_name_);
  RtlInitUnicodeString(&name, symbolic_link_name_);
  status = IoCreateSymbolicLink(&name, &device_name);
  if (!NT_SUCCESS(status)) {
    ERR("IoCreateSymbolicLink fail, symbolic_link_name=%ws, device_name=%ws, status=%#X\n", symbolic_link_name_, device_name_, status);
    return status;
  }
  INFO("IoCreateSymbolicLink Success, symbolic_link_name=%ws, device_name=%ws\n", symbolic_link_name_, device_name_);
  
  LEAVE();
  return status;

}

NTSTATUS Device::AddPendingIrp(PIRP irp) {
  NTSTATUS status;
  PendingIrp *pending_irp;

  PAGED_CODE();
  ENTER();
  if (!irp) {
    ERR("Device::AddPendingIrp irp is NULL\n");
    return STATUS_INVALID_PARAMETER_1;
  }
  if (!device_object_) {
    ERR("device_object_ is NULL\n");
    return STATUS_INVALID_DEVICE_STATE;
  }
  
  pending_irp = new PendingIrp;
  if (!pending_irp) {
    ERR("Device::AddPendingIrp new for PendingIrp is fail\n");
    status = STATUS_INSUFFICIENT_RESOURCES;
  } else {
    pending_irp->irp = irp;
    InsertTailList(&pending_irp_list_.node, &pending_irp->node);
    status = STATUS_SUCCESS;
    INFO("Device::AddPendingIrp Insert pending_irp=%p, irp=%p\n", pending_irp, pending_irp->irp);
  }

  LEAVE();
  return status;
}

NTSTATUS Device::CompleteAllPendingIrp() {
  NTSTATUS status;
  PLIST_ENTRY list_entry;
  PendingIrp *pending_irp;
  PIO_STACK_LOCATION stack_location;
  PVOID read_buffer;
  ULONG read_buffer_size;
  SIZE_T required_buffer_size;
  ULONG irp_count;
  LARGE_INTEGER interval;

  PAGED_CODE();
  ENTER();
  if (!device_object_) {
    ERR("device_object_ is NULL\n");
    return STATUS_INVALID_DEVICE_STATE;
  }
  irp_count = 0;
  while (!IsListEmpty(&pending_irp_list_.node)) {
    list_entry = RemoveHeadList(&pending_irp_list_.node);
    pending_irp = CONTAINING_RECORD(list_entry, PendingIrp, node);
    stack_location = IoGetCurrentIrpStackLocation(pending_irp->irp);
    read_buffer_size = stack_location->Parameters.Read.Length;
    read_buffer = pending_irp->irp->AssociatedIrp.SystemBuffer;
    INFO("%d pending_irp=%p, read_buffer=%p, read_buffer_size=%d\n", irp_count++, pending_irp, read_buffer, read_buffer_size);

    required_buffer_size = sizeof(WCHAR) * 24;
    if (read_buffer_size < required_buffer_size) {
      if (sizeof(UCHAR) <= read_buffer_size) {
        ((PUCHAR)read_buffer)[0] = (UCHAR)required_buffer_size;
      }
      pending_irp->irp->IoStatus.Status = STATUS_BUFFER_TOO_SMALL;
      pending_irp->irp->IoStatus.Information = 0;
    } else {
      interval.QuadPart = -1 * 1000 * 1000 * 10;
      KeDelayExecutionThread(KernelMode, FALSE, &interval);
      status = Time::GetCurrentTimeString(reinterpret_cast<PWSTR>(read_buffer));
      if (!NT_SUCCESS(status)) {
        ERR("Time::GetCurrentTime fail, read_buffer=%p, read_buffer_size=%d, status=%#X\n",
            read_buffer, read_buffer_size, status);
        pending_irp->irp->IoStatus.Status = status;
        pending_irp->irp->IoStatus.Information = 0;
      } else {
        CHAR current_time[24];
        SIZE_T time_size = sizeof(current_time);
        String::ConvertWcharToChar((PCWSTR)read_buffer, current_time, time_size);
        INFO("Current time:%s\n", current_time);
        pending_irp->irp->IoStatus.Status = STATUS_SUCCESS;
        pending_irp->irp->IoStatus.Information = required_buffer_size;
      }
    }
    IoCompleteRequest(pending_irp->irp, IO_NO_INCREMENT);
    INFO("Device::CompleteAllPendingIrp IoCompleteRequest pending_irp=%p, irp=%p\n", 
      pending_irp, pending_irp->irp);
    delete pending_irp;
    pending_irp = NULL;
  }
  LEAVE();
  return STATUS_SUCCESS;
}

NTSTATUS Device::FormatDeviceName(LPCWSTR device_name,
                                  LPWSTR standard_device_name, SIZE_T &size) {
  NTSTATUS status;
  PCWSTR prefix = L"\\Device\\";
  SIZE_T required_size;
  
  PAGED_CODE();
  ENTER();
  
  if (!device_name || !wcslen(device_name)) 
    return STATUS_INVALID_PARAMETER_1;
  if (!standard_device_name) 
    return STATUS_INVALID_PARAMETER_2;
  if (!size) 
    return STATUS_INVALID_PARAMETER_3;

  if (0 == wcsncmp(device_name, prefix, wcslen(prefix))) {
    required_size = sizeof(WCHAR) * (wcslen(device_name) + (sizeof(UNICODE_NULL) / sizeof(WCHAR)));
    if (required_size < size) {
      wcsncpy_s(standard_device_name, size, device_name, wcslen(device_name));  
      status = STATUS_SUCCESS;
    } else {
      status = STATUS_BUFFER_TOO_SMALL;
    }
    size = required_size;
  } else {
    required_size = sizeof(WCHAR) * (wcslen(prefix) + wcslen(device_name) + (sizeof(UNICODE_NULL) / sizeof(WCHAR)));
    if (required_size < size) {
      wcsncpy_s(standard_device_name, size, prefix, wcslen(prefix)); 
      wcsncat_s(standard_device_name, size, device_name, wcslen(device_name));
      status = STATUS_SUCCESS;
    } else {
      status = STATUS_BUFFER_TOO_SMALL;
    }
    size = required_size;
  }
  LEAVE();
  return status;
}

NTSTATUS Device::FormatSymbolicLinkName(LPCWSTR symbolic_link_name,
                                        LPWSTR standard_symbolic_link_name,
                                        SIZE_T &size) {
  NTSTATUS status;
  PCWSTR prefix = L"\\DosDevices\\";
  SIZE_T required_size;
  
  PAGED_CODE();
  ENTER();

  if (!symbolic_link_name || !wcslen(symbolic_link_name)) 
    return STATUS_INVALID_PARAMETER_1;
  if (!standard_symbolic_link_name) 
    return STATUS_INVALID_PARAMETER_2;
  if (!size) 
    return STATUS_INVALID_PARAMETER_3;
  
  if (0 == wcsncmp(symbolic_link_name, prefix, wcslen(prefix))) {
    required_size = sizeof(WCHAR) * (wcslen(symbolic_link_name) + (sizeof(UNICODE_NULL) / sizeof(WCHAR)));
    if (required_size < size) {
      RtlCopyMemory(standard_symbolic_link_name, symbolic_link_name, required_size);
      status = STATUS_SUCCESS;
    } else {
      status = STATUS_BUFFER_TOO_SMALL;
    }
    size = required_size;
  } else {
    required_size = sizeof(WCHAR) * (wcslen(prefix) + wcslen(symbolic_link_name) + (sizeof(UNICODE_NULL) / sizeof(WCHAR)));
    if (required_size < size) {
      wcsncpy_s(standard_symbolic_link_name, size, prefix, wcslen(prefix));
      wcsncat_s(standard_symbolic_link_name, size, symbolic_link_name, wcslen(symbolic_link_name));
      status = STATUS_SUCCESS;
    } else {
      status = STATUS_BUFFER_TOO_SMALL;
    }
    size = required_size;
  }
  LEAVE();
  return status;
}
